"""
WebSocket服务
"""
import json
import asyncio
from typing import List, Set
from fastapi import WebSocket, WebSocketDisconnect
from backend.utils.logger import log_manager
from backend.models.response_models import LogEntry

class WebSocketManager:
    """WebSocket连接管理器"""
    
    def __init__(self):
        self.active_connections: Set[WebSocket] = set()
        self.last_log_count = 0
        
    async def connect(self, websocket: WebSocket):
        """接受WebSocket连接"""
        await websocket.accept()
        self.active_connections.add(websocket)
        
        # 发送现有的日志
        await self.send_existing_logs(websocket)
        
        # 开始监听新日志
        asyncio.create_task(self.monitor_logs())
    
    def disconnect(self, websocket: WebSocket):
        """断开WebSocket连接"""
        self.active_connections.discard(websocket)
    
    async def send_existing_logs(self, websocket: WebSocket):
        """发送现有的日志给新连接的客户端"""
        try:
            logs = log_manager.get_logs(limit=100)  # 发送最近100条日志
            
            for log_entry in reversed(logs):  # 按时间顺序发送
                message = {
                    "type": "log",
                    "data": {
                        "id": log_entry.id,
                        "timestamp": log_entry.timestamp.isoformat(),
                        "level": log_entry.level,
                        "operation": log_entry.operation,
                        "message": log_entry.message,
                        "success": log_entry.success,
                        "details": log_entry.details
                    }
                }
                await websocket.send_text(json.dumps(message))
                
        except Exception as e:
            print(f"发送现有日志失败: {e}")
    
    async def broadcast_log(self, log_entry: LogEntry):
        """广播日志给所有连接的客户端"""
        if not self.active_connections:
            return
        
        message = {
            "type": "log",
            "data": {
                "id": log_entry.id,
                "timestamp": log_entry.timestamp.isoformat(),
                "level": log_entry.level,
                "operation": log_entry.operation,
                "message": log_entry.message,
                "success": log_entry.success,
                "details": log_entry.details
            }
        }
        
        message_str = json.dumps(message)
        
        # 使用副本避免在迭代过程中修改集合
        connections_copy = self.active_connections.copy()
        
        for websocket in connections_copy:
            try:
                await websocket.send_text(message_str)
            except Exception as e:
                print(f"发送日志到WebSocket失败: {e}")
                # 移除失效的连接
                self.active_connections.discard(websocket)
    
    async def monitor_logs(self):
        """监控新日志并推送"""
        while True:
            try:
                current_log_count = len(log_manager.logs)
                
                # 检查是否有新日志
                if current_log_count > self.last_log_count:
                    # 获取新日志
                    new_logs = log_manager.logs[self.last_log_count:]
                    
                    # 广播新日志
                    for log_entry in new_logs:
                        await self.broadcast_log(log_entry)
                    
                    self.last_log_count = current_log_count
                
                # 等待一段时间后再检查
                await asyncio.sleep(1)
                
            except Exception as e:
                print(f"监控日志失败: {e}")
                await asyncio.sleep(5)  # 出错后等待更长时间
    
    async def send_system_message(self, message: str, level: str = "INFO"):
        """发送系统消息"""
        if not self.active_connections:
            return
        
        system_message = {
            "type": "system",
            "data": {
                "timestamp": log_manager.logs[-1].timestamp.isoformat() if log_manager.logs else None,
                "level": level,
                "message": message
            }
        }
        
        message_str = json.dumps(system_message)
        
        connections_copy = self.active_connections.copy()
        
        for websocket in connections_copy:
            try:
                await websocket.send_text(message_str)
            except Exception as e:
                print(f"发送系统消息失败: {e}")
                self.active_connections.discard(websocket)
    
    async def send_connection_status(self, connected: bool, bootstrap_servers: str = None):
        """发送连接状态更新"""
        if not self.active_connections:
            return
        
        status_message = {
            "type": "connection_status",
            "data": {
                "connected": connected,
                "bootstrap_servers": bootstrap_servers,
                "timestamp": log_manager.logs[-1].timestamp.isoformat() if log_manager.logs else None
            }
        }
        
        message_str = json.dumps(status_message)
        
        connections_copy = self.active_connections.copy()
        
        for websocket in connections_copy:
            try:
                await websocket.send_text(message_str)
            except Exception as e:
                print(f"发送连接状态失败: {e}")
                self.active_connections.discard(websocket)

# 全局WebSocket管理器实例
websocket_manager = WebSocketManager() 